'''
python 3.10
    pip install Pillow
    pip install pymupdf
    pip install opencv-python
    pip install  paddlepaddle paddleocr
    pip install paddleclas
    (可能需要，一般不需要安装)pip install paddledet
'''

'''
PaddlePaddle version: 2.6.2
PaddleOCR version: 2.9.1
'''

import fitz
import os
import logging
from datetime import datetime
import shutil
import cv2
import numpy as np
from paddleocr import PPStructure
import traceback

class PDFProcessor:
    def __init__(self, languages=['ch', 'en'], use_gpu=True):
        self._setup_logging()

        # 初始化设备
        self.use_gpu = use_gpu

        # 初始化版面分析模型
        self.table_engine = PPStructure(
            show_log=True,
            image_orientation=True,
            use_gpu=use_gpu,
            lang=languages[0],
            layout=True,
            table=True,
            ocr=True,
            recovery=True
        )

    def _setup_logging(self):
        """设置日志"""
        log_dir = 'logs'
        os.makedirs(log_dir, exist_ok=True)
        log_file = os.path.join(log_dir, f'pdf_process_{datetime.now().strftime("%Y%m%d_%H%M%S")}.log')

        logging.basicConfig(
            level=logging.INFO,
            format='%(asctime)s - %(levelname)s - %(message)s',
            handlers=[
                logging.FileHandler(log_file, encoding='utf-8'),
                logging.StreamHandler()
            ]
        )
        self.logger = logging.getLogger(__name__)

    def _get_pdf_name(self, pdf_path):
        """获取PDF文件名（不含扩展名）"""
        base_name = os.path.basename(pdf_path)
        return os.path.splitext(base_name)[0]

    def analyze_layout(self, image_path):
        """使用PP-Structure进行版面分析"""
        try:
            # 读取图像
            img = cv2.imread(image_path)
            result = self.table_engine(img)

            # 初始化布局信息字典
            layout_info = {
                'text_regions': [],
                'title_regions': [],
                'table_regions': [],
                'figure_regions': [],
                'list_regions': []
            }

            # 处理每个检测到的区域
            for region in result:
                # 确保region是字典类型
                if not isinstance(region, dict):
                    continue

                region_type = region.get('type', '')
                box = region.get('bbox', None)

                if not box:
                    continue

                region_info = {
                    'box': box,
                    'text': '',
                    'confidence': 0.0
                }

                # 提取文本内容
                if 'res' in region:
                    if isinstance(region['res'], dict):
                        region_info['text'] = region['res'].get('text', '')
                    elif isinstance(region['res'], list):
                        # 如果res是列表，合并所有文本
                        texts = []
                        for item in region['res']:
                            if isinstance(item, dict):
                                text = item.get('text', '')
                                if text:
                                    texts.append(text)
                        region_info['text'] = '\n'.join(texts)

                # 根据类型存储区域信息
                if region_type == 'text':
                    layout_info['text_regions'].append(region_info)
                elif region_type == 'title':
                    layout_info['title_regions'].append(region_info)
                elif region_type == 'table':
                    layout_info['table_regions'].append(region_info)
                elif region_type == 'figure':
                    layout_info['figure_regions'].append(region_info)
                elif region_type == 'list':
                    layout_info['list_regions'].append(region_info)

            return layout_info

        except Exception as e:
            self.logger.error(f"版面分析出错: {str(e)}")
            self.logger.error(traceback.format_exc())
            return None

    def visualize_layout(self, image_path, layout_info, output_path):
        """可视化版面分析结果"""
        try:
            # 读取原始图像
            image = cv2.imread(image_path)
            if image is None:
                self.logger.error(f"无法读取图像: {image_path}")
                return False

            # 定义不同类型区域的颜色 (BGR格式)
            colors = {
                'text_regions': (0, 255, 0),  # 绿色
                'title_regions': (255, 0, 0),  # 蓝色
                'table_regions': (0, 0, 255),  # 红色
                'figure_regions': (255, 255, 0),  # 青色
                'list_regions': (255, 0, 255)  # 粉色
            }

            # 为每种类型绘制边界框
            for region_type, regions in layout_info.items():
                if region_type not in colors:
                    continue

                color = colors[region_type]
                for region in regions:
                    if 'box' not in region:
                        continue

                    box = region['box']
                    # 确保坐标为整数
                    x1, y1 = int(box[0]), int(box[1])
                    x2, y2 = int(box[2]), int(box[3])

                    # 绘制矩形框
                    cv2.rectangle(image, (x1, y1), (x2, y2), color, 2)

                    # 添加类型标签
                    label = region_type.replace('_regions', '')
                    font = cv2.FONT_HERSHEY_SIMPLEX
                    font_scale = 0.5
                    thickness = 1

                    # 获取文本大小
                    (text_width, text_height), _ = cv2.getTextSize(label, font, font_scale, thickness)

                    # 绘制文本背景
                    cv2.rectangle(image,
                                (x1, y1 - text_height - 5),
                                (x1 + text_width, y1),
                                color,
                                -1)  # -1 表示填充矩形

                    # 绘制文本
                    cv2.putText(image,
                              label,
                              (x1, y1 - 5),
                              font,
                              font_scale,
                              (255, 255, 255),  # 白色文本
                              thickness)

            # 保存结果
            cv2.imwrite(output_path, image)
            self.logger.info(f"已保存可视化结果到: {output_path}")
            return True

        except Exception as e:
            self.logger.error(f"可视化版面分析结果时出错: {str(e)}")
            self.logger.error(traceback.format_exc())
            return False

    def _process_single_page(self, page, page_number, pdf_dir, dpi=300):
        """处理单个页面（使用版面分析）"""
        try:
            # 创建必要的目录
            pdfimg_dir = os.path.join(pdf_dir, 'pdfimg')
            vis_dir = os.path.join(pdf_dir, 'visualization')
            texts_dir = os.path.join(pdf_dir, 'texts')

            for dir_path in [pdfimg_dir, vis_dir, texts_dir]:
                os.makedirs(dir_path, exist_ok=True)

            # 保存页面图像
            image_file = os.path.join(pdfimg_dir, f'page_{page_number + 1}.png')
            pix = page.get_pixmap(dpi=dpi)
            pix.save(image_file)

            # 进行版面分析
            layout_info = self.analyze_layout(image_file)

            if layout_info:
                # 生成可视化结果
                vis_output = os.path.join(vis_dir, f'page_{page_number + 1}_layout.png')
                self.visualize_layout(image_file, layout_info, vis_output)

                # 用于收集该页面的所有文本
                page_text = []

                # 处理不同类型的区域
                for region_type, regions in layout_info.items():
                    if not regions:
                        continue

                    try:
                        # 创建区域类型对应的目录
                        region_dir = os.path.join(pdf_dir, region_type)
                        os.makedirs(region_dir, exist_ok=True)

                        for idx, region in enumerate(regions):
                            try:
                                # 确保box存在且格式正确
                                if 'box' not in region or len(region['box']) != 4:
                                    continue

                                box = region['box']
                                img = cv2.imread(image_file)
                                if img is None:
                                    continue

                                # 确保坐标为整数且在有效范围内
                                y1, y2 = max(0, int(box[1])), min(img.shape[0], int(box[3]))
                                x1, x2 = max(0, int(box[0])), min(img.shape[1], int(box[2]))

                                if x2 <= x1 or y2 <= y1:
                                    continue

                                # 裁剪区域图像
                                cropped = img[y1:y2, x1:x2]

                                # 保存区域图像
                                image_path = os.path.join(region_dir, f'page_{page_number + 1}_{idx + 1}.png')
                                cv2.imwrite(image_path, cropped)

                                # 处理文本内容
                                if region.get('text'):
                                    # 添加区域类型标记和文本到页面文本集合
                                    text_content = f"[{region_type.replace('_regions', '').upper()}]\n{region['text']}\n"
                                    page_text.append(text_content)

                                    # 保存单独的文本文件
                                    text_path = os.path.join(region_dir, f'page_{page_number + 1}_{idx + 1}.txt')
                                    with open(text_path, 'w', encoding='utf-8') as f:
                                        f.write(region['text'])

                            except Exception as e:
                                self.logger.error(f'处理区域 {idx + 1} 时出错: {str(e)}')
                                continue

                    except Exception as e:
                        self.logger.error(f'处理区域类型 {region_type} 时出错: {str(e)}')
                        continue

                # 保存整页文本
                if page_text:
                    page_text_path = os.path.join(texts_dir, f'page_{page_number + 1}.txt')
                    with open(page_text_path, 'w', encoding='utf-8') as f:
                        f.write(f"=== 第 {page_number + 1} 页 ===\n\n")
                        f.write('\n'.join(page_text))

            self.logger.info(f'页面 {page_number + 1} 处理完成')
            return True

        except Exception as e:
            self.logger.error(f'处理页面 {page_number + 1} 时出错: {str(e)}')
            self.logger.error(traceback.format_exc())
            return False

    def _process_pdf(self, pdf_path, pdf_dir):
        """统一处理PDF文件"""
        try:
            # 创建必要的目录
            for dir_name in ['pdfimg', 'visualization', 'texts']:
                os.makedirs(os.path.join(pdf_dir, dir_name), exist_ok=True)

            pdf_document = fitz.open(pdf_path)
            total_pages = pdf_document.page_count

            # 处理每一页
            for page_number in range(total_pages):
                page = pdf_document[page_number]
                self._process_single_page(page, page_number, pdf_dir)

            # 合并所有文本文件
            self._merge_text_files(pdf_dir, total_pages)
            pdf_document.close()

            self.logger.info("PDF处理完成")

        except Exception as e:
            self.logger.error(f"处理PDF时出错: {str(e)}")
            self.logger.error(traceback.format_exc())

    def _merge_text_files(self, pdf_dir, total_pages):
        """合并所有文本文件"""
        merged_file = os.path.join(pdf_dir, 'merged_result.txt')
        texts_dir = os.path.join(pdf_dir, 'texts')

        with open(merged_file, 'w', encoding='utf-8') as outfile:
            for page_num in range(total_pages):
                page_file = os.path.join(texts_dir, f'page_{page_num + 1}.txt')
                if os.path.exists(page_file):
                    with open(page_file, 'r', encoding='utf-8') as infile:
                        outfile.write(f'\n=== 第 {page_num + 1} 页 ===\n\n')
                        outfile.write(infile.read())
                        outfile.write('\n\n')

    def process_pdf(self, pdf_path, output_dir="output"):
        """处理PDF文件的主入口"""
        try:
            # 创建输出目录
            os.makedirs(output_dir, exist_ok=True)

            # 获取PDF名称并创建带时间戳的专属目录
            pdf_name = self._get_pdf_name(pdf_path)
            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
            pdf_dir_name = f"{pdf_name}_{timestamp}"
            pdf_dir = os.path.join(output_dir, pdf_dir_name)

            # 如果目录已存在，先删除（这种情况不太可能发生，因为有时间戳）
            if os.path.exists(pdf_dir):
                shutil.rmtree(pdf_dir)
            os.makedirs(pdf_dir)

            # 统一处理PDF
            self._process_pdf(pdf_path, pdf_dir)

            self.logger.info(f"PDF处理完成，输出目录: {pdf_dir}")

        except Exception as e:
            self.logger.error(f"处理PDF时出错: {str(e)}")
            self.logger.error(traceback.format_exc())


if __name__ == "__main__":
    # 使用示例
    processor = PDFProcessor(
        languages=['ch', 'en'],
        use_gpu=False  # 如果有GPU可用，设置为True
    )

    pdf_path = "./input/重大气象信息专报第8期(1)_纯图版.pdf"
    processor.process_pdf(
        pdf_path,
        output_dir="output"
    )